#include <Dynamics/Pdm/pdm_cpu.h>
#include <Utilities/SpaceHash/space_hash_cpu.h>

using namespace std;

namespace PhysLeo {

template<typename T>
PdmCpu<T>::PdmCpu():
    ForceModule<T>(),
    pdm_material_type_(PdmMaterialType::LINEAR), 
    homogeneous_(true), 
    pdm_invertible_(true), 
    pdm_weight_type_(PdmWeightType::AVERAGE)
{
}

template<typename T>
void PdmCpu<T>::allocateResource(SimData<T>& sim_data)
{
    //buffer data
    ptr_rest_pos_ = sim_data.ptr_rest_pos_;
    ptr_pos_ = sim_data.ptr_pos_;
    ptr_force_ = sim_data.ptr_force_;
    ptr_vert_volume_ = sim_data.ptr_vert_volume_;
    ptr_vert_material_ = sim_data.ptr_vert_material_;

    //property
    horizon_radius_ = sim_data.horizon_radius_;
    homogeneous_ = sim_data.homogeneous_;
    pdm_material_type_ = sim_data.pdm_material_type_;
    pdm_invertible_ = sim_data.pdm_invertible_;
    pdm_weight_type_ = sim_data.pdm_weight_type_;

    int vert_num = ptr_vert_volume_->size();

    //allocate space for neighbor list
    ptr_neighbor_list_ = make_shared<BufferData<NeighborList>>(vert_num);
    ptr_neighbor_list_->allocateCpu();

    SpaceHashCpu<T> space_hash;
    auto bbox = space_hash.bbox(ptr_rest_pos_);
    space_hash.setSpace(horizon_radius_, bbox[0], bbox[1], 2 * SPACE_HASH_MAX_NEIGHBOR);
    space_hash.queryNeighbor(ptr_rest_pos_, ptr_neighbor_list_, horizon_radius_, horizon_radius_*0.1f);

    //allocate space for weight
    ptr_weight_ = make_shared<BufferData<PdmWeight<T>>>(vert_num);
    ptr_weight_->allocateCpu();
    initWeights();

    //allocate sapce for reference matrix Dm
    if(pdm_invertible_)
    {
        ptr_dm_inverse_ = make_shared<BufferData<glm::tmat3x3<T>>>(vert_num);
        ptr_dm_inverse_->allocateCpu();
        constructDm();
    }

    //allocate space for volume dilation
    ptr_volume_dilation_ = make_shared<BufferData<T>>(vert_num);
    ptr_volume_dilation_->allocateCpu();
}

template<typename T>
T PdmCpu<T>::deformEnergy(std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_pos)
{
    volumeDilation(ptr_pos);

    auto rest_pos_ptr = ptr_rest_pos_->ptrCpu();
    auto pos_ptr = ptr_pos->ptrCpu();
    const auto neighborlist_ptr = ptr_neighbor_list_->ptrCpu();
    auto weight_ptr = ptr_weight_->ptrCpu();
    auto volume_dilation_ptr = ptr_volume_dilation_->ptrCpu();
    auto material_ptr = ptr_vert_material_->ptrCpu();
    auto volume_ptr = ptr_vert_volume_->ptrCpu();

    const int vert_num = ptr_rest_pos_->size();

    const auto material_param_num = PdmMaterial<T>::numMaterialParams(pdm_material_type_);

    T total_energy = 0.0f;
    for (auto vert_id = 0; vert_id < vert_num; ++vert_id)
    {
        auto temp_material_ptr = material_ptr + (homogeneous_ ? 0 : vert_id * material_param_num);

        const auto neighbor_size = neighborlist_ptr[vert_id].size_;

        T energy = 0.0f;

        if(pdm_invertible_)
        {
            auto dm_inverse_ptr = ptr_dm_inverse_->ptrCpu();
            // construct current shape matrix Ds
            glm::tmat3x3<T> ds(0.0f);
            for(auto i=0;i<neighbor_size;++i)
            {
                auto dst_index = neighborlist_ptr[vert_id].ids_[i];
                auto vm = rest_pos_ptr[dst_index] - rest_pos_ptr[vert_id];
                auto vs = pos_ptr[dst_index] - pos_ptr[vert_id];
                ds += weight_ptr[vert_id][i] * glm::outerProduct(vs, vm);
            }
            //compute deformation gradient F
            auto F = ds * dm_inverse_ptr[vert_id];

            for(auto i=0;i<neighbor_size;++i)
            {
                auto dst_index = neighborlist_ptr[vert_id].ids_[i];
                glm::tvec3<T> origin = rest_pos_ptr[dst_index] - rest_pos_ptr[vert_id];
                glm::tvec3<T> deform = F * origin;
                energy += volume_ptr[vert_id] * weight_ptr[vert_id][i] * PdmMaterial<T>::energy(origin, deform, volume_dilation_ptr[vert_id], temp_material_ptr, pdm_material_type_);
            }
        }
        else
        {
            for (auto i = 0; i < neighbor_size; ++i)
            {
                auto dst_index = neighborlist_ptr[vert_id].ids_[i];
                glm::tvec3<T> origin = rest_pos_ptr[dst_index] - rest_pos_ptr[vert_id];
                glm::tvec3<T> deform = pos_ptr[dst_index] - pos_ptr[vert_id];
                energy += volume_ptr[vert_id] * weight_ptr[vert_id][i] * PdmMaterial<T>::energy(origin, deform, volume_dilation_ptr[vert_id], temp_material_ptr, pdm_material_type_);
            }
        }
        total_energy += energy;
    }
    return total_energy;
}

template<typename T>
void PdmCpu<T>::internalForce(std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_pos, std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_force)
{
    volumeDilation(ptr_pos);

    auto rest_pos_ptr = ptr_rest_pos_->ptrCpu();
    auto pos_ptr = ptr_pos->ptrCpu();
    auto force_ptr = ptr_force->ptrCpu();
    const auto neighborlist_ptr = ptr_neighbor_list_->ptrCpu();
    auto weight_ptr = ptr_weight_->ptrCpu();
    auto volume_dilation_ptr = ptr_volume_dilation_->ptrCpu();
    auto material_ptr = ptr_vert_material_->ptrCpu();
    auto volume_ptr = ptr_vert_volume_->ptrCpu();

    const int vert_num = ptr_rest_pos_->size();

    const auto material_param_num = PdmMaterial<T>::numMaterialParams(pdm_material_type_);

    for (auto vert_id = 0; vert_id < vert_num; ++vert_id)
    {
        auto temp_material_ptr = material_ptr + (homogeneous_ ? 0 : vert_id * material_param_num);

        const int neighbor_size = neighborlist_ptr[vert_id].size_;

        glm::tvec3<T> center_force(0.0f);

        if(pdm_invertible_)
        {
            auto dm_inverse_ptr = ptr_dm_inverse_->ptrCpu();
            // construct current shape matrix Ds
            glm::tmat3x3<T> ds(0.0f);
            for (auto i = 0; i<neighbor_size; ++i)
            {
                auto dst_index = neighborlist_ptr[vert_id].ids_[i];
                auto vm = rest_pos_ptr[dst_index] - rest_pos_ptr[vert_id];
                auto vs = pos_ptr[dst_index] - pos_ptr[vert_id];
                ds += weight_ptr[vert_id][i] * glm::outerProduct(vs, vm);
            }
            //compute deformation gradient F
            auto F = ds * dm_inverse_ptr[vert_id];

            for (auto i = 0; i < neighbor_size; ++i)
            {
                auto dst_index = neighborlist_ptr[vert_id].ids_[i];
                glm::tvec3<T> origin = rest_pos_ptr[dst_index] - rest_pos_ptr[vert_id];
                glm::tvec3<T> deform = pos_ptr[dst_index] - pos_ptr[vert_id];
                auto temp_force = PdmMaterial<T>::force(origin, deform, volume_dilation_ptr[vert_id], temp_material_ptr, pdm_material_type_);
                //Y-|X|dir(FX)
                glm::tvec3<T> force = volume_ptr[vert_id] * weight_ptr[vert_id][i] * glm::length(temp_force)*glm::normalize(deform - glm::length(origin)*glm::normalize(F*origin));
                center_force += force;
                force_ptr[dst_index] += -force;
            }
        }
        else
        {
            for (auto i = 0; i < neighbor_size; ++i)
            {
                auto dst_index = neighborlist_ptr[vert_id].ids_[i];
                glm::tvec3<T> origin = rest_pos_ptr[dst_index] - rest_pos_ptr[vert_id];
                glm::tvec3<T> deform = pos_ptr[dst_index] - pos_ptr[vert_id];

                glm::tvec3<T> force = volume_ptr[vert_id] * weight_ptr[vert_id][i] * PdmMaterial<T>::force(origin, deform, volume_dilation_ptr[vert_id], temp_material_ptr, pdm_material_type_);
                center_force += force;
                force_ptr[dst_index] += -force;
            }
        }
        force_ptr[vert_id] += center_force;
    }
}

template<typename T>
void PdmCpu<T>::volumeDilation(std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_pos)
{
    auto rest_pos_ptr = ptr_rest_pos_->ptrCpu();
    auto pos_ptr = ptr_pos->ptrCpu();
    const auto neighborlist_ptr = ptr_neighbor_list_->ptrCpu();
    auto weight_ptr = ptr_weight_->ptrCpu();
    auto volume_dilation_ptr = ptr_volume_dilation_->ptrCpu();

    const int vert_num = ptr_rest_pos_->size();

    for (auto vert_id = 0; vert_id < vert_num; ++vert_id)
    {
        const int neighbor_size = neighborlist_ptr[vert_id].size_;
        T theta = 0.0f;
        for (auto i = 0; i < neighbor_size; ++i)
        {
            auto dst_index = neighborlist_ptr[vert_id].ids_[i];
            glm::tvec3<T> origin = rest_pos_ptr[dst_index] - rest_pos_ptr[vert_id];
            glm::tvec3<T> deform = pos_ptr[dst_index] - pos_ptr[vert_id];

            auto s = glm::length(deform) / glm::length(origin);
            theta += weight_ptr[vert_id][i] * s;
        }
        volume_dilation_ptr[vert_id] = theta;
    }
}

template<typename T>
void PdmCpu<T>::initWeights()
{
    auto rest_pos_ptr = ptr_rest_pos_->ptrCpu();
    const auto neighborlist_ptr = ptr_neighbor_list_->ptrCpu();
    auto vert_volume_ptr = ptr_vert_volume_->ptrCpu();
    auto weight_ptr = ptr_weight_->ptrCpu();

    const int vert_num = ptr_rest_pos_->size();
    for(auto i=0;i<vert_num;++i)
    {
        const auto neighbor_size = neighborlist_ptr[i].size_;
        T total_weight = 0.0f;
        for(auto j=0;j<neighbor_size;++j)
        {
            auto dst_index = neighborlist_ptr[i].ids_[j];
            auto original_length = glm::length(rest_pos_ptr[i] - rest_pos_ptr[dst_index]);
            switch (pdm_weight_type_)
            {
            case PdmWeightType::AVERAGE:
                weight_ptr[i][j] = 1.0f;
                break;
            case PdmWeightType::CENTER:
                weight_ptr[i][j] = horizon_radius_ / (horizon_radius_ + original_length);
                break;
            case PdmWeightType::OUTER:
                weight_ptr[i][j] = (horizon_radius_ + original_length) / horizon_radius_;
                break;
            default:
                break;
            }
            weight_ptr[i][j] *= vert_volume_ptr[dst_index];
            total_weight += weight_ptr[i][j];
        }
        for(auto j=0;j<neighbor_size;++j)
        {
            //normalization
            weight_ptr[i][j] /= total_weight;
        }
    }
}

template<typename T>
void PdmCpu<T>::constructDm()
{
    auto rest_pos_ptr = ptr_rest_pos_->ptrCpu();
    const auto neighborlist_ptr = ptr_neighbor_list_->ptrCpu();
    auto weight_ptr = ptr_weight_->ptrCpu();
    auto dm_inverse_ptr = ptr_dm_inverse_->ptrCpu();

    const int vert_num = ptr_rest_pos_->size();
    for (auto i = 0; i<vert_num; ++i)
    {
        const auto neighbor_size = neighborlist_ptr[i].size_;
        glm::tmat3x3<T> dm(0.0f);
        for (auto j = 0; j<neighbor_size; ++j)
        {
            auto dst_index = neighborlist_ptr[i].ids_[j];
            auto vm = rest_pos_ptr[dst_index] - rest_pos_ptr[i];
            dm += weight_ptr[i][j] * glm::outerProduct(vm, vm);
        }

        // this place need special carefulness, Dm may be a singular matrix, so it can not be easily inversed.
        glm::tmat3x3<T> eigen_vec;
        glm::tvec3<T> diag;
        glm::tmat3x3<T> sigma(0.0f);
        glm::eigenDecomposition(dm, eigen_vec, diag);
        for(int j=0;j<3;++j)
        {
            if (diag[j] < 0.001*horizon_radius_) sigma[j][j] = 1000.0f / horizon_radius_;
            else sigma[j][j] = 1.0f / diag[j];
        }
        dm_inverse_ptr[i] = eigen_vec * sigma*glm::transpose(eigen_vec);
    }
}

template class PdmCpu<float>;
template class PdmCpu<double>;

}